'''
实现股票池，条件是0 < PE_TTM < 30, 按照PE_TT正序排列，最多取100只股票;
再平衡周期为20个交易日
主要的方法包括：
stockPool:找到两个日期之间的所有出的票
findOutStocks: 找出当前被调出的股票
evaluateStockPool: 对股票池的性能做初步验证，确保股票池能够带来Alpha
'''

from pymongo import MongoClient, ASCENDING, DESCENDING
import pandas as pd
import matplotlib.pyplot as plt

db = MongoClient('mongodb://localhost:27017')['quant']

'''
股票池
:param beginDate: 开始日期
:param endDate: 结束日期
:return: tuple, 所有调整日，以及调整日和代码列表对应的dict
'''
def stockPool(beginDate, endDate):
    adjustDateCodesDict = dict()

    # 获取指定时间区间的所有交易日
    timeCursor = db.daily_k.find(
        {'code': '000001.SH', 'time': {'$gte': beginDate, '$lte': endDate}},
        sort = [('time', ASCENDING)],
        projection = {'time': True},
        batch_size = 1000
    )

    allDates = [x['time'] for x in timeCursor]

    # 上一期的所有股票代码
    lastPhaseCodes = []
    # 调整周期是20个交易日
    adjustInterval = 20
    # 所有的调整日
    allAdjustDates = []

    #在调整日调整股票池
    for index in range(0, len(allDates),adjustInterval):
        # 保存调整日
        adjustDate = allDates[index]
        allAdjustDates.append(adjustDate)

        # 查询出调整当日，0 < pe_ttm < 30, 且非停牌的股票
        # 最重要的一点是，按照pe_ttm正序排列，只取前100只
        dailyCursor = db.daily_k.find(
            {'time': adjustDate, 'code':{'$regex': '[6|0|3]0\d{4}$'}, 'pe_ttm': {'$lt': 30, '$gt': 0}, 'trade_status': '交易'},
            sort = [('pe_ttm', ASCENDING)],
            projection = {'code': True},
            limit = 100
        )

        codes = [x['code'] for x in dailyCursor]

        # 本期股票列表
        thisPhaseCodes = []

        # 查询出上次股票池中正在停牌的股票
        if len(lastPhaseCodes) > 0:
            suspensionCursor = db.daily_k.find(
                {'code': {'$in': lastPhaseCodes, 'time': adjustDate, 'trade_status': '停牌'}},
                projection = {'code': True}
            )
            suspensionCodes = [x['code'] for x in suspensionCursor]

            # 保留在股票池中正在停牌的股票
            thisPhaseCodes = suspensionCodes

        # 用新的股票将剩余的位置补齐
        thisPhaseCodes += codes[0 : 100 - len(thisPhaseCodes)]
        # 将本次股票设为下次运行时的上次股票池
        lastPhaseCodes = thisPhaseCodes

        # 建立该调整日和股票列表的对应关系
        adjustDateCodesDict[adjustDate] = thisPhaseCodes

    # 返回结果
    return allAdjustDates, adjustDateCodesDict

'''
找到上期入选本期被调出的股票，这些股票将必须卖出
:param last_phase_codes: 上期的股票列表
:param this_phase_codes: 本期的股票列表
:return: 被调出的股票列表
'''
def findOutStocks(lastPhaseCodes, thisPhaseCodes):
    outStocks = []

    for code in lastPhaseCodes:
        if code not in thisPhaseCodes:
            outStocks.append(code)

    return outStocks

'''
对股票池做一个简单的评价
'''
def evaluateStockPool():
    # 设定评测周期
    adjustDates, codesDict = stockPool('2015-01-01', '2015-12-31')

    # 用DataFrame保存收益
    dfProfit = pd.DataFrame(columns=['profit', 'hs300'])

    dfProfit.loc[adjustDates[0]] = {'profit': 0, 'hs300': 0}

    hs300BeginValue = db.daily_k.find_one({'code': '000300.SH',  'time' : adjustDates[0]})['close']

    # 通过净值计算累计收益
    netValue = 1
    for index in range(1, len(adjustDates) - 1):
        lastAdjustDate = adjustDates[index - 1]
        currentAdjustDate = adjustDates[index]
        # 获取上一期的股票池
        codes = codesDict[lastAdjustDate]

        # 构建股票代码和后复权买入价格的股票
        codeBuyCloseDict = dict()
        buyDailyCursor = db.daily_k_ba.find(
            {'code': {'$in': codes}, 'time': lastAdjustDate},
            projection = {'close': True, 'code': True}
        )

        for buyDaily in buyDailyCursor:
            code = buyDaily['code']
            codeBuyCloseDict[code] = buyDaily['close']

        # 获取到期的股价
        sellDailyCursor = db.daily_k_ba.find(
            {'code': {'$in': codes}, 'time': currentAdjustDate},
            projection = {'close': True, 'code': True}
        )

        # 计算单期收益
        profitSum = 0
        for sellDaily in sellDailyCursor:
            code = sellDaily['code']

            buyColse = codeBuyCloseDict[code]
            sellClose = sellDaily['close']

            profitSum += (sellClose - buyColse) / buyColse

        profit = round(profitSum/len(codes), 4)
        hs300Close = db.daily_k.find(
            {'code': '000300.SH', 'time': currentAdjustDate}
        )['close']

        # 计算净值和累积收益
        netValue = netValue * (1 + profit)
        dfProfit.loc[currentAdjustDate] = {
            'profit': round((netValue - 1) * 100, 4),
            'hs300': round((hs300Close - hs300BeginValue) * 100 / hs300BeginValue, 4)
        }

    # 绘制曲线
    dfProfit.plot(title='Stock Pool Evaluation Result', kind = 'line')
    # 显示图像
    plt.show()

if __name__ == '__main__':
    evaluateStockPool()